{
  "metadata": {
    "dataExplorerConfig": {},
    "accelerator": "GPU",
    "colab": {
      "background_execution": "on",
      "collapsed_sections": [],
      "machine_shape": "hm",
      "name": "Torchrec Planner Introduction.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "pytorch (cinder)",
      "name": "bento_kernel_pytorch_cinder",
      "language": "python",
      "metadata": {
        "kernel_name": "bento_kernel_pytorch_cinder",
        "nightly_builds": true,
        "fbpkg_supported": true,
        "cinder_runtime": true,
        "is_prebuilt": true
      }
    },
    "language_info": {
      "name": "python"
    },
    "last_server_session_id": "b4be4a2a-0467-4335-9dd6-415fcbcbd20b",
    "last_kernel_id": "9550fcbd-27ae-4709-a721-8740d1c32a26",
    "last_base_url": "https://devgpu013.atn5.facebook.com:8090/",
    "last_msg_id": "de1399ef-04b73c8cf8a0dd2e69231cce_4649",
    "outputWidgetContext": {}
  },
  "nbformat": 4,
  "nbformat_minor": 0,
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HWBOrwVSnrNE",
        "originalKey": "8a17ffef-2882-43c7-a1c6-746a04dee662",
        "showInput": false,
        "code_folding": [],
        "hidden_ranges": []
      },
      "source": [
        "## **Automated Planner Tutorial**\n",
        "The planner attempts to identify the optimal sharding plan by evaluating a series of proposals which are statically analyzed and fed into an integer partitioner.  The planner is able to automatically select near optimal sharding plans for a wide range of hardware setups, allowing users to scale performance seamlessly from their local development environment to large scale production hardware. \n",
        "\n",
        "The built-in planner is packaged with TorchRec, which can be directly accessed via `EmbeddingShardingPlanner` API or indirectly via `DistributedModelParallel` API.\n",
        "\n",
        "It is recommended to use the default planner as is, and most performance improvements are accomplished by modifying the default arguments via the public facing API.\n",
        "\n",
        "In this tutorial we will explore how to:\n",
        "\n",
        "- setup a `Topology`, run the planner to generate a `ShardingPlan`, and interpret basic diagnostic information\n",
        "- use `ParameterConstaints` to select sharding type and provide pooling factors\n",
        "- trigger `UVM Caching` when model exceeds GPU memory (happens automatically!)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7iY7Uv11mJYK",
        "originalKey": "cbc09bbd-2ff7-48a6-a70c-caa2270493d8",
        "showInput": false,
        "code_folding": [],
        "hidden_ranges": []
      },
      "source": [
        "### Initial Setup (if needed)\n",
        "\n",
        "\n",
        "Installing TorchRec will also install [FBGEMM](https://github.com/pytorch/fbgemm), a collection of CUDA kernels and GPU enabled operations to run "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tUnIw-ZREQJy",
        "originalKey": "39ef497d-ff26-49bd-8db5-e37265fb9ba4",
        "code_folding": [],
        "hidden_ranges": []
      },
      "source": [
        "# install torchrec\n",
        "!pip3 install torchrec-nightly"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b6EHgotRXFQh",
        "originalKey": "4504f378-37f5-4f61-94fe-fef5fd3c5c47"
      },
      "source": [
        "The following steps are needed for the Colab runtime to detect the added shared libraries. The runtime searches for shared libraries in /usr/lib, so we copy over the libraries which were installed in /usr/local/lib/. **This is a very necessary step, only in the Colab runtime**. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_P45pDteRcWj",
        "originalKey": "117f5483-ff58-4d79-a813-db8e2581f894"
      },
      "source": [
        "!sudo cp /usr/local/lib/lib* /usr/lib/"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "n5_X2WOAYG3c",
        "originalKey": "8a171d9c-9aca-4437-822c-8326aa91c0d4"
      },
      "source": [
        "\\**Restart your runtime at this point for the newly installed packages to be seen.** Run the step below immediately after restarting so that python knows where to look for packages. **Always run this step after restarting the runtime.**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8cktNrh8R9rC",
        "originalKey": "d32b32d4-8b57-4888-8297-9ed158fa325b",
        "collapsed": false,
        "requestMsgId": "c0d336eb-d5e6-481b-bebc-099b257809e5",
        "customOutput": null,
        "executionStartTime": 1655358135143,
        "executionStopTime": 1655358135253
      },
      "source": [
        "import sys\n",
        "sys.path = ['', '/env/python', '/usr/local/lib/python37.zip', '/usr/local/lib/python3.7', '/usr/local/lib/python3.7/lib-dynload', '/usr/local/lib/python3.7/site-packages', './.local/lib/python3.7/site-packages']"
      ],
      "execution_count": 7,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZdSUWBRxoP8R",
        "originalKey": "cdca0ca7-5f11-4f0d-93dc-6bda0f984e86",
        "showInput": false,
        "code_folding": [],
        "hidden_ranges": []
      },
      "source": [
        "### Constructing Embedding Tables\n",
        "\n",
        "To showcase the versatility of the planner, we are going start by creating create 2 embedding bags, a small one and a huge one."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "e7UQBuG09hbj",
        "originalKey": "ff74ef44-04a3-4440-a923-cfece3b1927f",
        "collapsed": false,
        "requestMsgId": "941e03bd-a120-4f35-95f1-1948585a3aa3",
        "customOutput": null,
        "executionStartTime": 1655485006707,
        "executionStopTime": 1655485006749,
        "code_folding": [],
        "hidden_ranges": []
      },
      "source": [
        "import torchrec\n",
        "\n",
        "small_table = torchrec.EmbeddingBagConfig(\n",
        "  name=\"small_table\",\n",
        "  embedding_dim=64,\n",
        "  num_embeddings=10000000,\n",
        "  feature_names=[\"small_table_feature\"],\n",
        "  pooling=torchrec.PoolingType.SUM,\n",
        ")\n",
        "\n",
        "huge_table = torchrec.EmbeddingBagConfig(\n",
        "  name=\"huge_table\",\n",
        "  embedding_dim=1024,\n",
        "  num_embeddings=50000000,\n",
        "  feature_names=[\"large_table_feature\"],\n",
        "  pooling=torchrec.PoolingType.SUM,\n",
        ")"
      ],
      "execution_count": 230,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "originalKey": "20e859b3-2e5d-4080-9959-529ad310f115",
        "showInput": false,
        "customInput": null,
        "code_folding": [],
        "hidden_ranges": []
      },
      "source": [
        "###  Running the standalone Planner\n",
        "TorchRec seperates model sharding into two separate stages:\n",
        "- Planning stage: Determine \"how\" to shard the model for a given sharder(s) and a run-time enviornment (`Topology`).  Output is known as the Sharding Plan\n",
        "- Sharding stage: Use the given sharder(s) to shard the model in accordance to the Sharding Plan.  This needs executed in the run-time environment.\n",
        "\n",
        "This separation of responsiblities allows the planning step to be executed independent of the run-time environment, and we can leverage that to explore various hypothetical scenarios in this tutorial.\n",
        "\n",
        "Ok, lets start with the small table, and look at the planner output on a single gpu run-time."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "originalKey": "d829ac75-e968-42f5-8352-5deba39e2a0d",
        "showInput": true,
        "customInput": null,
        "collapsed": false,
        "requestMsgId": "caaad682-f7a3-4c07-a00f-22a83251ca23",
        "customOutput": null,
        "executionStartTime": 1655483520076,
        "executionStopTime": 1655483520257,
        "code_folding": [],
        "hidden_ranges": []
      },
      "source": [
        "from torchrec.distributed.planner import EmbeddingShardingPlanner, Topology\n",
        "from torchrec.distributed.embeddingbag import (\n",
        "    EmbeddingBagCollectionSharder,\n",
        ")\n",
        "\n",
        "# Toy model, We will allocate on device \"meta\", to avoid allocation of data pointer\n",
        "small_toy_model = torchrec.EmbeddingBagCollection(\n",
        "    device=\"meta\", \n",
        "    tables=[small_table]\n",
        ")\n",
        "\n",
        "# Create a run-time environment, 1 gpu\n",
        "topology_1gpu = Topology(\n",
        "    world_size=1, \n",
        "    compute_device=\"cuda\"\n",
        ")\n",
        "\n",
        "planner = EmbeddingShardingPlanner(\n",
        "    topology=topology_1gpu,\n",
        "    debug=True,\n",
        ")\n",
        "\n",
        "sharding_plan = planner.plan(small_toy_model, sharders=[EmbeddingBagCollectionSharder()])"
      ],
      "execution_count": 193,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 093200.096 stats.py:190] ###################################################################################################\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 093200.096 stats.py:192] #                                   --- Planner Statistics ---                                    #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 093200.096 stats.py:199] #           --- Evalulated 16 proposal(s), found 16 possible plan(s), ran for 0.02s ---           #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 093200.097 stats.py:202] # ----------------------------------------------------------------------------------------------- #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 093200.098 stats.py:205] #      Rank     HBM (GB)     DDR (GB)     Perf (ms)     Input (MB)     Output (MB)     Shards     #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 093200.098 stats.py:205] #    ------   ----------   ----------   -----------   ------------   -------------   --------     #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 093200.099 stats.py:205] #         0     2.4 (7%)     0.0 (0%)         0.001            0.0            0.12      TW: 1     #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 093200.099 stats.py:207] #                                                                                                 #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 093200.100 stats.py:209] # Input: MB/iteration, Output: MB/iteration, Shards: number of tables                             #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 093200.101 stats.py:211] # HBM: est. peak memory usage for shards - parameter, comms, optimizer, and gradients             #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 093200.101 stats.py:212] #                                                                                                 #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 093200.102 stats.py:218] # Compute Kernels:                                                                                #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 093200.102 stats.py:220] #   batched_fused: 1                                                                              #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 093200.103 stats.py:223] #                                                                                                 #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 093200.103 stats.py:224] # Parameter Info:                                                                                 #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 093200.104 stats.py:227] #             FQN     Sharding     Compute Kernel     Perf (ms)     Ranks                         #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 093200.104 stats.py:227] #           -----   ----------   ----------------   -----------   -------                         #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 093200.105 stats.py:227] #    .small_table           TW      batched_fused         0.001         0                         #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 093200.105 stats.py:229] ###################################################################################################\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "originalKey": "eb7e2cf5-0f7e-4a3c-8e96-6b5555c6aeb1",
        "showInput": false,
        "customInput": null,
        "code_folding": [],
        "hidden_ranges": []
      },
      "source": [
        "The `Planner Statistics` gives a high level summary of sharding plan, some interesting things to note:\n",
        "- A total of 16 different proposals were evaluated, all of which were valid plans.   The Planner will return the best sharding plan, defined as the plan with the lowest `Perf` value\n",
        "- In this situation, the table is `Table-wise` (TW) sharded, using `batched_fused` kernel and is estimated to consume 2.4 GB of HBM memory \n",
        "\n",
        "Now lets look at what happens when we increase world_size to 2:\n",
        "\n",
        ""
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "originalKey": "c851a77d-2f91-4ba5-bdad-27d9462005a6",
        "showInput": true,
        "customInput": null,
        "code_folding": [],
        "hidden_ranges": [],
        "collapsed": false,
        "requestMsgId": "eeba9011-c77e-4ed4-9775-23d7566fb49c",
        "customOutput": null,
        "executionStartTime": 1655483522576,
        "executionStopTime": 1655483522800
      },
      "source": [
        "# Create a run-time environment, 2 gpu\n",
        "topology_2gpu = Topology(\n",
        "    world_size=2, \n",
        "    compute_device=\"cuda\"\n",
        ")\n",
        "\n",
        "planner = EmbeddingShardingPlanner(\n",
        "    topology=topology_2gpu,\n",
        "    debug=True,\n",
        ")\n",
        "\n",
        "sharding_plan = planner.plan(small_toy_model, sharders=[EmbeddingBagCollectionSharder()])"
      ],
      "execution_count": 194,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 093202.598 stats.py:190] ###################################################################################################\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 093202.598 stats.py:192] #                                   --- Planner Statistics ---                                    #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 093202.599 stats.py:199] #           --- Evalulated 16 proposal(s), found 16 possible plan(s), ran for 0.02s ---           #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 093202.600 stats.py:202] # ----------------------------------------------------------------------------------------------- #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 093202.601 stats.py:205] #      Rank     HBM (GB)     DDR (GB)     Perf (ms)     Input (MB)     Output (MB)     Shards     #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 093202.602 stats.py:205] #    ------   ----------   ----------   -----------   ------------   -------------   --------     #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 093202.603 stats.py:205] #         0     1.2 (4%)     0.0 (0%)         0.003            0.0            0.12      RW: 1     #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 093202.603 stats.py:205] #         1     1.2 (4%)     0.0 (0%)         0.003            0.0            0.12      RW: 1     #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 093202.604 stats.py:207] #                                                                                                 #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 093202.605 stats.py:209] # Input: MB/iteration, Output: MB/iteration, Shards: number of tables                             #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 093202.605 stats.py:211] # HBM: est. peak memory usage for shards - parameter, comms, optimizer, and gradients             #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 093202.606 stats.py:212] #                                                                                                 #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 093202.606 stats.py:218] # Compute Kernels:                                                                                #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 093202.607 stats.py:220] #   batched_fused: 1                                                                              #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 093202.608 stats.py:223] #                                                                                                 #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 093202.608 stats.py:224] # Parameter Info:                                                                                 #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 093202.609 stats.py:227] #             FQN     Sharding     Compute Kernel     Perf (ms)     Ranks                         #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 093202.610 stats.py:227] #           -----   ----------   ----------------   -----------   -------                         #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 093202.610 stats.py:227] #    .small_table           RW      batched_fused         0.005       0-1                         #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 093202.611 stats.py:229] ###################################################################################################\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "originalKey": "29b1f26e-7b2c-4ec8-aac3-787a476a2ead",
        "showInput": false,
        "customInput": null,
        "code_folding": [],
        "hidden_ranges": []
      },
      "source": [
        "With two GPUs:\n",
        "- The same parameter is now sharded `Row-wise` (RW).\n",
        "- Each rank's estimated peak memory usage is 1.2 GB, reflecting the parameter is split evenly between two ranks.\n",
        "\n",
        "One question at this someone may have is, how do I know this is the best plan?  For example would how does this compare to just using Table-wise.   Fortunately, the planner API supports this by utilizing `ParameterConstraints`.\n",
        ""
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "originalKey": "ebed64b2-36ff-4ce8-8cee-80cce8c38e45",
        "showInput": true,
        "customInput": null,
        "code_folding": [],
        "hidden_ranges": [],
        "collapsed": false,
        "requestMsgId": "46dcf72f-434f-4f8e-b1fc-e2e6c5cfac82",
        "customOutput": null,
        "executionStartTime": 1655484067106,
        "executionStopTime": 1655484067272
      },
      "source": [
        "from torchrec.distributed.planner.types import ParameterConstraints\n",
        "\n",
        "constraints = {\n",
        "    \"small_table\": ParameterConstraints(\n",
        "        sharding_types=[ShardingType.TABLE_WISE.value]\n",
        "    ),\n",
        "}\n",
        "\n",
        "planner = EmbeddingShardingPlanner(\n",
        "    topology=topology_2gpu,\n",
        "    constraints=constraints,\n",
        "    debug=True,\n",
        ")\n",
        "plan = planner.plan(small_toy_model, sharders=[EmbeddingBagCollectionSharder()])"
      ],
      "execution_count": 210,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 094107.148 stats.py:190] ###################################################################################################\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 094107.148 stats.py:192] #                                   --- Planner Statistics ---                                    #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 094107.149 stats.py:199] #            --- Evalulated 3 proposal(s), found 3 possible plan(s), ran for 0.00s ---            #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 094107.149 stats.py:202] # ----------------------------------------------------------------------------------------------- #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 094107.150 stats.py:205] #      Rank     HBM (GB)     DDR (GB)     Perf (ms)     Input (MB)     Output (MB)     Shards     #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 094107.151 stats.py:205] #    ------   ----------   ----------   -----------   ------------   -------------   --------     #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 094107.151 stats.py:205] #         0     2.4 (7%)     0.0 (0%)         0.003           0.01            0.25      TW: 1     #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 094107.152 stats.py:205] #         1     0.0 (0%)     0.0 (0%)           0.0            0.0             0.0      TW: 0     #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 094107.152 stats.py:207] #                                                                                                 #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 094107.153 stats.py:209] # Input: MB/iteration, Output: MB/iteration, Shards: number of tables                             #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 094107.154 stats.py:211] # HBM: est. peak memory usage for shards - parameter, comms, optimizer, and gradients             #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 094107.154 stats.py:212] #                                                                                                 #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 094107.155 stats.py:218] # Compute Kernels:                                                                                #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 094107.155 stats.py:220] #   batched_fused: 1                                                                              #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 094107.156 stats.py:223] #                                                                                                 #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 094107.156 stats.py:224] # Parameter Info:                                                                                 #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 094107.157 stats.py:227] #             FQN     Sharding     Compute Kernel     Perf (ms)     Ranks                         #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 094107.157 stats.py:227] #           -----   ----------   ----------------   -----------   -------                         #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 094107.158 stats.py:227] #    .small_table           TW      batched_fused         0.003         0                         #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 094107.158 stats.py:229] ###################################################################################################\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "originalKey": "9b421bb2-9017-4e32-b403-a15f63db7e7a",
        "showInput": false,
        "customInput": null,
        "code_folding": [],
        "hidden_ranges": []
      },
      "source": [
        "Requiring the planner to use Table-wise results in nearly identical `Perf`, with Rank 0 holding the parameter\n",
        "\n",
        "Another important use of `ParameterConstraints` is to provide the planner with information regarding the pooling factor(s) for each table.  In models with a high number of tables this becomes the dominant factor in estimating model `Perf` and critial to planner's ability to accurately estimate peak HBM memory usage.\n",
        ""
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "originalKey": "282c8299-5d46-47dc-b6d8-06d9edf8d8d9",
        "showInput": true,
        "customInput": null,
        "code_folding": [],
        "hidden_ranges": [],
        "collapsed": false,
        "requestMsgId": "08cd6a11-47fe-4a78-b2ae-af53b6cf49fe",
        "customOutput": null,
        "executionStartTime": 1655484847030,
        "executionStopTime": 1655484847198
      },
      "source": [
        "constraints = {\n",
        "    \"small_table\": ParameterConstraints(\n",
        "        pooling_factors=[200.0],\n",
        "    ),\n",
        "}\n",
        "\n",
        "planner = EmbeddingShardingPlanner(\n",
        "    topology=topology_2gpu,\n",
        "    constraints=constraints,\n",
        "    debug=True,\n",
        ")\n",
        "plan = planner.plan(small_toy_model, sharders=[EmbeddingBagCollectionSharder()])"
      ],
      "execution_count": 222,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 095407.058 stats.py:190] ###################################################################################################\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 095407.059 stats.py:192] #                                   --- Planner Statistics ---                                    #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 095407.060 stats.py:199] #           --- Evalulated 16 proposal(s), found 16 possible plan(s), ran for 0.02s ---           #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 095407.060 stats.py:202] # ----------------------------------------------------------------------------------------------- #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 095407.061 stats.py:205] #      Rank     HBM (GB)     DDR (GB)     Perf (ms)     Input (MB)     Output (MB)     Shards     #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 095407.061 stats.py:205] #    ------   ----------   ----------   -----------   ------------   -------------   --------     #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 095407.062 stats.py:205] #         0     1.2 (4%)     0.0 (0%)         0.087           0.78            0.12      RW: 1     #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 095407.063 stats.py:205] #         1     1.2 (4%)     0.0 (0%)         0.087           0.78            0.12      RW: 1     #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 095407.063 stats.py:207] #                                                                                                 #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 095407.064 stats.py:209] # Input: MB/iteration, Output: MB/iteration, Shards: number of tables                             #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 095407.065 stats.py:211] # HBM: est. peak memory usage for shards - parameter, comms, optimizer, and gradients             #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 095407.066 stats.py:212] #                                                                                                 #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 095407.067 stats.py:218] # Compute Kernels:                                                                                #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 095407.067 stats.py:220] #   batched_fused: 1                                                                              #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 095407.068 stats.py:223] #                                                                                                 #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 095407.069 stats.py:224] # Parameter Info:                                                                                 #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 095407.069 stats.py:227] #             FQN     Sharding     Compute Kernel     Perf (ms)     Ranks                         #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 095407.070 stats.py:227] #           -----   ----------   ----------------   -----------   -------                         #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 095407.071 stats.py:227] #    .small_table           RW      batched_fused         0.173       0-1                         #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 095407.072 stats.py:229] ###################################################################################################\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "originalKey": "33a61e5f-4c0c-4b24-93dc-d3528a0e9d30",
        "showInput": false,
        "customInput": null,
        "code_folding": [],
        "hidden_ranges": []
      },
      "source": [
        "In this case, we can see by providing a pooling factor of 200.0, the planner has `Pref` has increased to 0.087 ms from 0.003 ms.\n",
        "\n",
        "Finally, let's look quickly at what happens when we have a huge table that will not fit in HBM memory."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "originalKey": "a937478c-8274-4796-a9c0-a8f6b9eb835a",
        "showInput": true,
        "customInput": null,
        "code_folding": [],
        "hidden_ranges": [],
        "collapsed": false,
        "requestMsgId": "b6be4de1-9581-468a-a124-a2a422d73d4b",
        "customOutput": null,
        "executionStartTime": 1655485011580,
        "executionStopTime": 1655485011733
      },
      "source": [
        "from torchrec.distributed.types import ShardingType\n",
        "\n",
        "\n",
        "# Toy model\n",
        "huge_toy_model = torchrec.EmbeddingBagCollection(device=\"meta\", tables=[huge_table])\n",
        "\n",
        "planner = EmbeddingShardingPlanner(\n",
        "    topology=topology_2gpu,\n",
        "    debug=True,\n",
        ")\n",
        "sharding_plan = planner.plan(huge_toy_model, sharders=[EmbeddingBagCollectionSharder()])"
      ],
      "execution_count": 231,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 095651.613 stats.py:190] ###################################################################################################\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 095651.614 stats.py:192] #                                   --- Planner Statistics ---                                    #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 095651.614 stats.py:199] #           --- Evalulated 16 proposal(s), found 8 possible plan(s), ran for 0.02s ---            #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 095651.615 stats.py:202] # ----------------------------------------------------------------------------------------------- #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 095651.617 stats.py:205] #      Rank     HBM (GB)     DDR (GB)     Perf (ms)     Input (MB)     Output (MB)     Shards     #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 095651.617 stats.py:205] #    ------   ----------   ----------   -----------   ------------   -------------   --------     #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 095651.619 stats.py:205] #         0   19.1 (60%)   95.4 (75%)         0.543           0.03             2.0      CW: 4     #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 095651.619 stats.py:205] #         1   19.1 (60%)   95.4 (75%)         0.543           0.03             2.0      CW: 4     #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 095651.620 stats.py:207] #                                                                                                 #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 095651.621 stats.py:209] # Input: MB/iteration, Output: MB/iteration, Shards: number of tables                             #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 095651.622 stats.py:211] # HBM: est. peak memory usage for shards - parameter, comms, optimizer, and gradients             #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 095651.622 stats.py:212] #                                                                                                 #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 095651.623 stats.py:218] # Compute Kernels:                                                                                #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 095651.624 stats.py:220] #   batched_fused_uvm_caching: 1                                                                  #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 095651.625 stats.py:223] #                                                                                                 #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 095651.626 stats.py:224] # Parameter Info:                                                                                 #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 095651.627 stats.py:227] #            FQN     Sharding              Compute Kernel     Perf (ms)             Ranks         #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 095651.628 stats.py:227] #          -----   ----------            ----------------   -----------           -------         #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 095651.628 stats.py:227] #    .huge_table           CW   batched_fused_uvm_caching         1.086   0,0,0,0,1,1,1,1         #\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "I0617 095651.630 stats.py:229] ###################################################################################################\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "originalKey": "0e32b349-54e1-4531-99af-fa90dd2b715d",
        "showInput": false,
        "customInput": null,
        "code_folding": [],
        "hidden_ranges": []
      },
      "source": [
        "The planner recognizes that table cannot fit in HBM device memory, and selects `UVM_CACHING` to accomdate the large parameter.  Also, planner had determined `Column-wise` (CW) sharding is optimal in this situation.\n",
        "\n",
        "I hope this tutorial has given you a brief introduction into the planner.  More detailed functionality is outlined in the planner API documentation.\n",
        "\n",
        "\n",
        ""
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "originalKey": "d0adae77-4e79-4199-8e74-b7faf783cff2",
        "showInput": true,
        "customInput": null
      },
      "source": [
        ""
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}
